package gibbs

import (
	"math/rand"
    "cs466/findmotif/motif"
)

func ProbabilityGivenPWM(p motif.PWM, site []byte) (prob float64) {
	prob = 1

	for mi, nucleotide := range site {
		prob *= p[mi][motif.NucleotideToBits[nucleotide]]
	}

	return
}

func CandidateProbabilities(p motif.PWM, seq []byte) (probs []float64) {
	l, ml := len(seq), len(p)

	probs = make([]float64, l - ml + 1)

	for i, _ := range probs {
		probs[i] = ProbabilityGivenPWM(p, seq[i:i+ml])
	}

	return
}

func WeightedSelect(w []float64) int {
	cdf := make([]float64, len(w))

	cdf[0] = 0
	for i := range cdf[1:] {
		cdf[i+1] = cdf[i] + w[i]
	}

	sum := cdf[len(cdf)-1] + w[len(w)-1]

	for i := range cdf {
		cdf[i] /= sum
	}

    return CDFSelect(cdf, rand.Float64())
}

func CDFSelect(cdf []float64, x float64) int {
	// modified binary search to find the largest element less than x
	var s []float64
	for s = cdf; len(s) > 1; {
		m := len(s)/2
		switch {
			case x >= s[m]:
				s = s[m:]
			case x < s[m-1]:
				s = s[:m-1]
			default:
				// x is between s[m-1] and s[m], so s[m-1] is the largest element less than x
				s = s[m-1:m]
		}
	}

	// if cdf and s are backed by the same memory region (which they are), this
	// is the index in cdf of the first element of s.
	return cap(cdf) - cap(s)
}

func PWMFor(sites [][]byte) motif.PWM {
	ml := len(sites[0])
	sc := float64(len(sites))

    ret := make(motif.PWM, ml)
    for i := range ret {
        ret[i] = make(motif.Column, 4)
    }

	for _, site := range sites {
		for i, nucleotide := range site {
			ret[i][motif.NucleotideToBits[nucleotide]] += 1/sc
		}
	}

    return ret
}

